`include "defines.v"

module if_stage(
    input  wire                             clk,
    input  wire                             rst,

    input  wire                             stall_i,

    input  wire                             transfer_flag_i,
    input  wire [`RAM_ADDR_WIDTH - 1 : 0]   transfer_addr_i,

    input  wire                             exception_flag_i,
    input  wire [`RAM_ADDR_WIDTH - 1 : 0]   exception_addr_i,

    output wire [`RAM_ADDR_WIDTH - 1 : 0]   iram_addr,
    input  wire [`IRAM_DATA_WIDTH - 1 : 0]  iram_rdata,
    output wire                             iram_valid,
    input  wire                             iram_ready,

    output wire                             stall_pc_flag_o,    
    output wire                             flush_if_id_flag_o,

    output wire [`RAM_ADDR_WIDTH - 1 : 0]   inst_addr,
    output wire [`INST_WIDTH - 1: 0]        inst
);
    reg  [`RAM_ADDR_WIDTH - 1 : 0] pc;
    
    wire [`RAM_ADDR_WIDTH - 1 : 0] pc_add_4;
    wire [`RAM_ADDR_WIDTH - 1 : 0] pc_value;
    
    reg fetch_state;

    reg                             transfer_flag;
    reg [`RAM_ADDR_WIDTH - 1 : 0]   transfer_addr;
    wire                            transfer_flag_value;
    wire[`RAM_ADDR_WIDTH - 1 : 0]   transfer_addr_value;

    reg                             exception_flag;
    reg [`RAM_ADDR_WIDTH - 1 : 0]   exception_addr;
    wire                            exception_flag_value;
    wire[`RAM_ADDR_WIDTH - 1 : 0]   exception_addr_value;

    assign transfer_flag_value  = fetch_state == 1'b0    && iram_ready == 1'b0 &&  transfer_flag_i == 1'b1 ? 1'b1 :
                                  transfer_flag == 1'b1  && iram_ready == 1'b1                             ? 1'b0 : transfer_flag;
    assign exception_flag_value = fetch_state == 1'b0    && iram_ready == 1'b0 && exception_flag_i == 1'b1 ? 1'b1 :
                                  exception_flag == 1'b1 && iram_ready == 1'b1                             ? 1'b0 : exception_flag;

    assign transfer_addr_value  =  transfer_flag == 1'b0 &&  transfer_flag_i == 1'b1 ?  transfer_addr_i : transfer_addr;
    assign exception_addr_value = exception_flag == 1'b0 && exception_flag_i == 1'b1 ? exception_addr_i : exception_addr;
    
    assign pc_add_4 = pc + 4;
    assign pc_value = exception_flag    ? exception_addr   :
                      exception_flag_i  ? exception_addr_i :
                      transfer_flag     ? transfer_addr    :
                      transfer_flag_i   ? transfer_addr_i  :
                                          pc_add_4;

    always @(posedge clk) begin
        if(rst == 1'b1) begin
            transfer_flag  <= 0;
            transfer_addr  <= 0;

            exception_flag <= 0;
            exception_addr <= 0;
        end
        else begin
            transfer_flag  <= transfer_flag_value;
            transfer_addr  <= transfer_addr_value;

            exception_flag <= exception_flag_value;
            exception_addr <= exception_addr_value;
        end 
    end
    // fetch an instruction
    always@(posedge clk)
    begin
        if(rst == 1'b1) begin
            pc <= `PC_START;
        end
        else begin
            if(stall_i == 1'b1) begin
                pc <= pc;
            end
            else begin
                pc <= pc_value;
            end 
        end
    end
    always @(posedge clk) begin
        if(rst == 1'b1) begin
            fetch_state <= 0;
        end 
        else begin
            case (fetch_state)
                1'b0: begin
                    if(iram_ready == 1'b0 && (transfer_flag_i | exception_flag_i) == 1'b1) begin
                        fetch_state <= 1'b1;
                    end
                end
                1'b1: begin
                    if(iram_ready) begin
                        fetch_state <= 1'b0;
                    end
                end
                default:;
            endcase
        end
    end
    //
    //
    // to axi
    assign iram_addr  = pc;
    assign iram_valid = rst == 1'b1 ? 0 : 1;
    //
    //
    // to ctrl
    assign stall_pc_flag_o = ~iram_ready;
    assign flush_if_id_flag_o = fetch_state == 1'b0 && iram_ready == 1'b1 && (transfer_flag_i | exception_flag_i) == 1'b0 ? 1'b0 : 1'b1;
    //
    //
    //to if_id
    assign inst      = pc[2] == 1'b1 ? iram_rdata[63:32] : iram_rdata[31: 0];
    // assign inst = iram_rdata;
    assign inst_addr = pc;
endmodule

